/**
 * @file Spectrum.h
 * @author martin_z_he
 * @brief mostly taken from PBRT, this is shared between CPP & CUDA
 * @version 0.1
 * @date 2021-11-18
 *
 * @copyright Copyright (c) 2021
 *
 */

#pragma once
#ifndef OBR_SPECTRUM_H
#define OBR_SPECTRUM_H

#include "helper_math.h"

#ifndef __CUDA_ARCH__
#include <stdio.h>
#endif

namespace OBR
{

// const float spectrumWavelengths[] = {440, 530, 700}; // current wavelengths
const float spectrumWavelengths[] = {550};

/**
 * @brief sampled spectrum to represent color
 *
 * @tparam nCoefficients
 */
template <int nCoefficients> class CoefficientSpectrum
{
public:
    // CoefficientSpectrum Public Methods
    __host__ __device__ CoefficientSpectrum(float v = 0.0f)
    {
        for (int i = 0; i < nCoefficients; ++i)
            c[i] = v;
    }
    __host__ __device__ CoefficientSpectrum(float *v)
    {
        if (v != nullptr)
            for (int i = 0; i < nCoefficients; ++i)
                c[i] = v[i];
    }
    __host__ __device__ CoefficientSpectrum(const CoefficientSpectrum &s)
    {
        for (int i = 0; i < nCoefficients; ++i)
            c[i] = s.c[i];
    }
    __host__ __device__ CoefficientSpectrum &operator=(const CoefficientSpectrum &s)
    {
        for (int i = 0; i < nCoefficients; ++i)
            c[i] = s.c[i];
        return *this;
    }
    __host__ __device__ void print() const
    {
        printf("[ ");
        for (int i = 0; i < nCoefficients; ++i)
        {
            printf("%f", c[i]);
            if (i != nCoefficients - 1)
                printf(", ");
        }
        printf("]");
    }
    __host__ __device__ CoefficientSpectrum &operator+=(const CoefficientSpectrum &s2)
    {
        for (int i = 0; i < nCoefficients; ++i)
            c[i] += s2.c[i];
        return *this;
    }
    __host__ __device__ CoefficientSpectrum operator+(const CoefficientSpectrum &s2) const
    {
        CoefficientSpectrum ret = *this;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] += s2.c[i];
        return ret;
    }
    __host__ __device__ CoefficientSpectrum operator-(const CoefficientSpectrum &s2) const
    {
        CoefficientSpectrum ret = *this;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] -= s2.c[i];
        return ret;
    }
    __host__ __device__ CoefficientSpectrum operator/(const CoefficientSpectrum &s2) const
    {
        CoefficientSpectrum ret = *this;
        for (int i = 0; i < nCoefficients; ++i)
        {
            // CHECK_NE(s2.c[i], 0);
            ret.c[i] /= s2.c[i];
        }
        return ret;
    }
    __host__ __device__ CoefficientSpectrum operator*(const CoefficientSpectrum &sp) const
    {
        CoefficientSpectrum ret = *this;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] *= sp.c[i];
        return ret;
    }
    __host__ __device__ CoefficientSpectrum &operator*=(const CoefficientSpectrum &sp)
    {
        for (int i = 0; i < nCoefficients; ++i)
            c[i] *= sp.c[i];
        return *this;
    }
    __host__ __device__ CoefficientSpectrum operator*(float a) const
    {
        CoefficientSpectrum ret = *this;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] *= a;
        return ret;
    }
    __host__ __device__ CoefficientSpectrum &operator*=(float a)
    {
        for (int i = 0; i < nCoefficients; ++i)
            c[i] *= a;
        return *this;
    }
    __host__ __device__ friend inline CoefficientSpectrum operator*(float a,
                                                                    const CoefficientSpectrum &s)
    {
        return s * a;
    }
    __host__ __device__ CoefficientSpectrum operator/(float a) const
    {
        CoefficientSpectrum ret = *this;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] /= a;
        return ret;
    }
    __host__ __device__ CoefficientSpectrum &operator/=(float a)
    {
        for (int i = 0; i < nCoefficients; ++i)
            c[i] /= a;
        return *this;
    }
    __host__ __device__ bool operator==(const CoefficientSpectrum &sp) const
    {
        for (int i = 0; i < nCoefficients; ++i)
            if (c[i] != sp.c[i])
                return false;
        return true;
    }
    __host__ __device__ bool operator!=(const CoefficientSpectrum &sp) const
    {
        return !(*this == sp);
    }
    __host__ __device__ bool isBlack() const
    {
        for (int i = 0; i < nCoefficients; ++i)
            if (c[i] != 0.)
                return false;
        return true;
    }
    __host__ __device__ bool isNan() const
    {
        for (int i = 0; i < nCoefficients; ++i)
            if (isnan(c[i]))
                return false;
        return true;
    }
    __host__ __device__ bool isInf() const
    {
        for (int i = 0; i < nCoefficients; ++i)
            if (isinf(c[i]))
                return false;
        return true;
    }

    __host__ __device__ bool isValid() const
    {
        return !(isNan() && isBlack() && isInf());
    }
    __host__ __device__ CoefficientSpectrum sqrt() const
    {
        CoefficientSpectrum ret;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] = sqrtf(c[i]);
        return ret;
    }
    __host__ __device__ CoefficientSpectrum pow(float e) const
    {
        CoefficientSpectrum ret;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] = powf(this->c[i], e);
        return ret;
    }
    __host__ __device__ CoefficientSpectrum operator-() const
    {
        CoefficientSpectrum ret;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] = -c[i];
        return ret;
    }
    __host__ __device__ CoefficientSpectrum exp() const
    {
        CoefficientSpectrum ret;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] = expf(c[i]);
        return ret;
    }
    __host__ __device__ CoefficientSpectrum clamp(float low = 0, float high = 1e10) const
    {
        CoefficientSpectrum ret;
        for (int i = 0; i < nCoefficients; ++i)
            ret.c[i] = fminf(high, fmaxf(low, c[i]));
        return ret;
    }
    __host__ __device__ float maxComponentValue() const
    {
        float m = c[0];
        for (int i = 1; i < nCoefficients; ++i)
            m = max(m, c[i]);
        return m;
    }
    __host__ __device__ float &operator[](int i)
    {
        return c[i];
    }
    __host__ __device__ void fromRGB(float r, float g, float b)
    {
        if constexpr (nCoefficients == 3)
        {
            c[0] = r / 255.0;
            c[1] = g / 255.0;
            c[2] = b / 255.0;
        }
        if constexpr (nCoefficients == 1)
        {
            c[0] = (r * 0.212671f + g * 0.715160f + b * 0.072169f) / 255.0;
        }
    }
    __host__ __device__ void toArray(float *dst) const
    {
        for (size_t i = 0; i < nCoefficients; i++)
        {
            dst[i] = c[i];
        }
    }

    // CoefficientSpectrum Public Data
    static const int nSamples = nCoefficients;

protected:
    // CoefficientSpectrum Protected Data
    float c[nCoefficients];
};

class RGBSpectrum : public CoefficientSpectrum<3>
{
public:
    __host__ __device__ RGBSpectrum(float v = 0.0f) : CoefficientSpectrum<3>(v)
    {
    }
    __host__ __device__ RGBSpectrum(float *v) : CoefficientSpectrum<3>(v)
    {
    }
    __host__ __device__ RGBSpectrum(float3 v) : CoefficientSpectrum<3>((float *)&v)
    {
    }
    __host__ __device__ RGBSpectrum(const RGBSpectrum &ref) : CoefficientSpectrum<3>(ref)
    {
    }
    __host__ __device__ RGBSpectrum(const CoefficientSpectrum<3> &ref) : CoefficientSpectrum<3>(ref)
    {
    }
    __host__ __device__ float3 toFloat3() const
    {
        return make_float3(c[0], c[1], c[2]);
    }
    __host__ __device__ float luminance() const
    {
        return c[0] * 0.212671f + c[1] * 0.715160f + c[2] * 0.072169f;
    }
};

class SingleWavelengthSpectrum : public CoefficientSpectrum<1>
{
public:
    __host__ __device__ SingleWavelengthSpectrum(float v = 0.0f) : CoefficientSpectrum<1>(v)
    {
    }
    __host__ __device__ SingleWavelengthSpectrum(float *v) : CoefficientSpectrum<1>(v)
    {
    }
    __host__ __device__ SingleWavelengthSpectrum(const SingleWavelengthSpectrum &ref)
        : CoefficientSpectrum<1>(ref)
    {
    }
    __host__ __device__ SingleWavelengthSpectrum(const RGBSpectrum &ref)
        : CoefficientSpectrum<1>(ref.luminance())
    {
    }
    __host__ __device__ SingleWavelengthSpectrum(const CoefficientSpectrum<1> &ref)
        : CoefficientSpectrum<1>(ref)
    {
    }
    __host__ __device__ float3 toFloat3() const
    {
        return make_float3(c[0], c[0], c[0]);
    }
    __host__ __device__ float luminance() const
    {
        return c[0];
    }
};

typedef SingleWavelengthSpectrum Spectrum;
// typedef RGBSpectrum Spectrum;

inline float rgb2luminance(const float rgb[3])
{
    return rgb[0] * 0.212671f + rgb[1] * 0.715160f + rgb[2] * 0.072169f;
}

inline void xyz2rgb(const float xyz[3], float rgb[3])
{
    rgb[0] = 3.240479f * xyz[0] - 1.537150f * xyz[1] - 0.498535f * xyz[2];
    rgb[1] = -0.969256f * xyz[0] + 1.875991f * xyz[1] + 0.041556f * xyz[2];
    rgb[2] = 0.055648f * xyz[0] - 0.204043f * xyz[1] + 1.057311f * xyz[2];
}

inline void rgb2xyz(const float rgb[3], float xyz[3])
{
    xyz[0] = 0.412453f * rgb[0] + 0.357580f * rgb[1] + 0.180423f * rgb[2];
    xyz[1] = 0.212671f * rgb[0] + 0.715160f * rgb[1] + 0.072169f * rgb[2];
    xyz[2] = 0.019334f * rgb[0] + 0.119193f * rgb[1] + 0.950227f * rgb[2];
}
} // namespace OBR

#endif